#include "ICP.h"
#include "kdTree.h"

double distance(const Point3d& a, const Point3d& b) {
	double x2 = (a.x-b.x) * (a.x-b.x);
	double y2 = (a.y-b.y) * (a.y-b.y);
	double z2 = (a.z-b.z) * (a.z-b.z);
	return std::sqrt(x2+y2+z2);
}

// return the index of the matched point in B
int FindClosedPoint(const vector<Point3d>& B, const Point3d& a) {
	int index = 0;
	double dis = distance(a, B[0]);
	for (int i=1; i<B.size(); i++) {
		double d = distance(a, B[i]);
		if (d < dis) {
			dis = d;
			index = i;
		}
	}
	return index;
}

void rotate_ensemble(const vector<Point3d>& B, vector<Point3d>& C, 
					 const Eigen::Matrix3d& R, const Eigen::Vector3d& t) {
	C.clear();
	for (int i=0; i<B.size(); i++) {
	    Vector3d b_original(B[i].x, B[i].y, B[i].z);
	    Vector3d c_transformed = R * b_original + t;
	    Point3d c_trans(c_transformed(0,0), c_transformed(1,0), c_transformed(2,0));
	    C.push_back(c_trans);
	}	
}

void pose_estimate_ICP(const vector<cv::Point3d>& A, 
					   vector<cv::Point3d>& B, 
					   Eigen::Matrix3d& R, Eigen::Vector3d& t) {
	// initialization
	R = Eigen::MatrixXd::Identity(3,3);
	t = Eigen::MatrixXd::Zero(3,1);

	double max_error = 1e-7;
	double average_error = 1;
	int loop = 1;
	int itermax = 200;
	
	// parameters
	int N = A.size();
	
	// containers
	vector<int> w(N);
	vector<int> matched_index(N);
	vector<Point3d> C(B.size());
	vector<double> error(N);
	double sum_error;
	
	while ( average_error > max_error && loop < itermax ) {
		cout << "loop: " << loop++ << endl;
		matched_index.clear();
		// build a kd Tree
		kdTree KD(B);
		sum_error = 0;
		for ( int i=0; i<N; i++ ) {
		    int *res = new int[1];
		    // int index = FindClosedPoint(B, A[i]);
		    
		    // kd search
		    
		    KD.findKNearests(A[i], 1, res);
		    int index = res[0];
		    
			matched_index.push_back(index);
			double d = distance(B[index], A[i]);
			error.push_back(d);
			sum_error += d;
		}
		average_error = sum_error / N;
		/*
		for (int i=0; i<N; i++) {
		    // error[i] > 2*average_error ? w.push_back(0) : w.push_back(1);
		    w.push_back(1);
		}
		*/

		// center of mass
		Point3d p1, p2;

    	for ( int i=0; i<N; i++ )
    	{
       		p1 += A[i];
        	p2 += B[matched_index[i]];
    	}
    	p1 = Point3d( Vec3d(p1) / N);
    	p2 = Point3d( Vec3d(p2) / N);
		
		// remove the center
		vector<Point3d> q1 (N), q2(N); 
    	for ( int i=0; i<N; i++ )
    	{
        	q1[i] = A[i] - p1;
        	q2[i] = B[matched_index[i]] - p2;
    	}
        
    	// compute q1*q2^T
    	Eigen::Matrix3d W = Eigen::Matrix3d::Zero();
    	for ( int i=0; i<N; i++ )
    	{
        	W += Eigen::Vector3d ( q1[i].x, q1[i].y, q1[i].z ) * Eigen::Vector3d ( q2[i].x, q2[i].y, q2[i].z ).transpose();
    	}
    	// cout<<"W="<<W<<endl;

    	// SVD on W
    	Eigen::JacobiSVD<Eigen::Matrix3d> svd ( W, Eigen::ComputeFullU|Eigen::ComputeFullV );
    	Eigen::Matrix3d U = svd.matrixU();
    	Eigen::Matrix3d V = svd.matrixV();
        
        /*
    	if (U.determinant() * V.determinant() < 0)
		{
        	for (int x = 0; x < 3; ++x)
        	{
            	U(x, 2) *= -1;
        	}
		}
		*/
    
    	// cout<<"U=\n"<<U<<endl;
    	// cout<<"V=\n"<<V<<endl;
        
        Eigen::Matrix3d d_R;
        Eigen::Vector3d d_t;
    	d_R = U* ( V.transpose() );
    	d_t = Eigen::Vector3d ( p1.x, p1.y, p1.z ) - d_R * Eigen::Vector3d ( p2.x, p2.y, p2.z );
        
        R = d_R * R;
        t = d_R * t + d_t;
        
        // update B
        rotate_ensemble(B, C, d_R, d_t);
        B.clear();
    	B.swap(C);
    	
		/*
    	cout << "dR = \n";
    	cout << d_R; cout << endl;
    	cout << "dt = "; cout << d_t.transpose(); cout << endl;
    	cout << "R = \n";
    	cout << R; cout << endl;
    	cout << "t = "; cout << t.transpose(); cout << endl;
    	*/
    	cout << "Average error = " << average_error << endl;
	}
							
}

